package com.gitee.feizns.explore.data_structure.tree.cus;

import java.util.LinkedList;
import java.util.Queue;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

/**
 * @author feizns
 * @since 2019/11/25
 */
public class TreeUtils {

    public static boolean isBalanced(Node root) {
        return root == null ? true : Math.abs(deep(root.left()) - deep(root.right())) < 2;
    }

    public static int deep(Node node) {
        int i = 0;
        if ( node != null ) {
            i++;
            int l = deep(node.left());
            int r = deep(node.right());
            i += (l > r ? l : r);
        }
        return i;
    }

    public static boolean isBst(Node root) {
        if ( root != null ) {
            if ( root.left() != null && root.left().key().compareTo(root.key()) > 0 )
                return false;
            if ( root.right() != null && root.right().key().compareTo(root.key()) < 0 )
                return false;
            return isBst(root.left()) && isBst(root.right());
        }
        return true;
    }

    public static <N extends Node, K extends Comparable<K>, V> void sequenceTraversal(Node<N, K, V> root, BiConsumer<K, V> consumer) {
        Queue<Node<N, K,V>> queue = new LinkedList<>();
        queue.add(root);
        sequenceTraversal(queue, consumer, queue::add);
    }

    public static <N extends Node, K extends Comparable<K>, V> void sequenceTraversal(Queue<Node<N, K, V>> queue, BiConsumer<K, V> consumer, Consumer<Node<N, K, V>> use) {
        while ( !queue.isEmpty() ) {
            Node<N, K, V> node = queue.remove();
            consumer.accept(node.key(), node.val());
            if ( node.left() != null )
                use.accept(node.left());
            if ( node.right() != null )
                use.accept(node.right());
        }
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N min(Node<N, K, V> root) {
        return root != null && root.left() == null ? (N) root : min(root.left());
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N rMin(N root) {
        return rMin(root, n -> {});
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N rMin(N root, Consumer<N> consumer) {
        if ( root != null && root.left() != null ) {
            root.left(rMin(root.left()));
            consumer.accept(root);
            return root;
        }
        return null;
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> PutRet<N, V> remove(N node, K k) {
//        PutRet<N, V> ret = null;
        if ( node != null ) {
            int comp = node.key().compareTo(k);
            V v = node.val();
            if ( comp > 0 ) {
                PutRet<N, V> tmp = remove(node.left(), k);
                node.left(tmp.node());
                return new PutRet<>(node, tmp.v());
            } else if ( comp < 0 ) {
                PutRet<N, V> tmp = remove(node.right(), k);
                node.right(tmp.node());
                return new PutRet<>(node, tmp.v());
            } else if ( comp == 0 ) {
                if ( node.left() != null && node.right() != null ) {
                    N min = TreeUtils.min(node.right());
                    N right = TreeUtils.rMin(node.right());
                    min.left(node.left());
                    min.right(right);
                    return new PutRet<>(min, v);
                } else if ( node.left() == null )
                    return new PutRet<>(node.right(), v);
                else if ( node.right() == null )
                    return new PutRet<>(node.left(), v);
            }
        }
        return new PutRet<>(node, null);
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N ll(N node) {
        return ll(node, n -> {});
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N ll(N node, Consumer<N> consumer) {
        N left = node.left();
        node.left(left.right());
        left.right(node);
        consumer.accept(node);
        consumer.accept(left);
        return left;
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N rr(N node) {
        return rr(node, n -> {});
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N rr(N node, Consumer<N> consumer) {
        N right = node.right();
        node.right(right.left());
        right.left(node);
        consumer.accept(node);
        consumer.accept(right);
        return right;
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N lr(N node) {
        return lr(node, n -> {});
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N lr(N node, Consumer<N> consumer) {
        N left = node.left();
        N lr = left.right();
        node.left(lr);
        left.right(lr.left());
        lr.left(left);
        consumer.accept(lr);
        consumer.accept(left);
        consumer.accept(node);
        return node;
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N lr2(N node, Consumer<N> consumer) {
        node.left(rr(node.left(), consumer));
        return ll(node, consumer);
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N rl(N node) {
        return rl(node, n -> {});
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N rl(N node, Consumer<N> consumer) {
        try {
            N right = node.right();
            N rl = right.left();
            node.right(rl);
            right.left(rl.right());
            rl.right(right);
            consumer.accept(rl);
            consumer.accept(right);
            consumer.accept(node);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return node;
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N rl2(N node, Consumer<N> consumer) {
        node.right(ll(node.right(), consumer));
        return rr(node, consumer);
    }

    public static <K extends Comparable<K>, V> void print(Node root) {
        Queue<Node> queue = new LinkedList<>();
        queue.add(root);
        print(queue);
    }

    public static void print(Queue<Node> queue) {
        Queue<Node> newQueue = new LinkedList<>();
        while ( !queue.isEmpty() ) {
            Node node = queue.remove();
            System.out.print(node.key() + " ");
            if ( node.left() != null )
                newQueue.add(node.left());
            if ( node.right() != null )
                newQueue.add(node.right());
        }
        System.out.println();
        if ( !newQueue.isEmpty() )
            print(newQueue);
    }

    public static void main(String[] args) {
        NodeImpl<Integer, Object> node = create();
        TreeUtils.print(node);
        node = rl2(node, n -> {});
        System.out.println("===");
        TreeUtils.print(node);
    }

    private static NodeImpl<Integer, Object> create() {
        NodeImpl<Integer, Object> node = new NodeImpl<>(4, null);
        NodeImpl<Integer, Object> node1 = new NodeImpl<>(6, null);
        NodeImpl<Integer, Object> node2 = new NodeImpl<>(5, null);
        node.right(node1);
        node1.left(node2);
        return node;
    }

    public static <N extends Node<N, K, V>, K extends Comparable<K>, V> N find(Node<N, K, V> root, K k) {
        if ( root != null ) {
            int weight = root.key().compareTo(k);
            if ( weight == 0 )
                return (N) root;
            else if ( weight > 0 )
                return find(root.left(), k);
            else
                return find(root.right(), k);
        }
        return null;
    }

    private static class NodeImpl<K extends Comparable<K>, V> implements Node<TreeUtils.NodeImpl, K, V> {

        K k;

        V v;

        NodeImpl<K, V> left, right;

        public NodeImpl(K k, V v) {
            this.k = k;
            this.v = v;
        }

        @Override
        public K key() {
            return k;
        }

        @Override
        public void key(K k) {
            this.k = k;
        }

        @Override
        public V val() {
            return v;
        }

        @Override
        public void val(V v) {
            this.v = v;
        }

        @Override
        public TreeUtils.NodeImpl left() {
            return left;
        }

        @Override
        public void left(TreeUtils.NodeImpl left) {
            this.left = left;
        }

        @Override
        public TreeUtils.NodeImpl right() {
            return right;
        }

        @Override
        public void right(TreeUtils.NodeImpl right) {
            this.right = right;
        }

    }

}
